package com.dc.schedule.spring;

import com.dc.schedule.api.model.ClientConfig;
import com.dc.schedule.api.model.StaticTaskRegistry;
import com.dc.schedule.api.model.StaticTaskRegistryConfig;
import com.dc.schedule.client.ScheduleClient;
import com.dc.schedule.client.socket.ScheduleSocketClient;
import com.dc.schedule.spring.anno.StaticTask;
import com.dc.schedule.spring.task.ContextualMethodInjectionTaskRunnable;
import com.dc.schedule.spring.task.ScheduleClientStaticTask;
import lombok.Getter;
import lombok.RequiredArgsConstructor;
import lombok.Setter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.aop.framework.AopInfrastructureBean;
import org.springframework.aop.framework.AopProxyUtils;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.DisposableBean;
import org.springframework.beans.factory.SmartInitializingSingleton;
import org.springframework.beans.factory.config.DestructionAwareBeanPostProcessor;
import org.springframework.beans.factory.support.MergedBeanDefinitionPostProcessor;
import org.springframework.beans.factory.support.RootBeanDefinition;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.context.ApplicationListener;
import org.springframework.context.EmbeddedValueResolverAware;
import org.springframework.context.event.ContextRefreshedEvent;
import org.springframework.core.MethodIntrospector;
import org.springframework.core.Ordered;
import org.springframework.core.annotation.AnnotationUtils;
import org.springframework.lang.NonNull;
import org.springframework.scheduling.TaskScheduler;
import org.springframework.scheduling.Trigger;
import org.springframework.scheduling.TriggerContext;
import org.springframework.scheduling.config.ScheduledTask;
import org.springframework.scheduling.config.ScheduledTaskHolder;
import org.springframework.scheduling.config.ScheduledTaskRegistrar;
import org.springframework.scheduling.support.CronTrigger;
import org.springframework.util.StringUtils;
import org.springframework.util.StringValueResolver;

import java.lang.annotation.Annotation;
import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Date;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.ScheduledExecutorService;
import java.util.stream.Collectors;

/**
 * 任务注解注册配置
 *
 * @see org.springframework.scheduling.annotation.ScheduledAnnotationBeanPostProcessor
 */
@Slf4j
@RequiredArgsConstructor
public class ScheduleClientBeanPostProcessor implements ScheduledTaskHolder, MergedBeanDefinitionPostProcessor, DestructionAwareBeanPostProcessor,
                                                                Ordered, EmbeddedValueResolverAware, ApplicationContextAware,
                                                                SmartInitializingSingleton, ApplicationListener<ContextRefreshedEvent>, DisposableBean {
    @Getter
    private final int order = LOWEST_PRECEDENCE;
    
    @Setter
    private StringValueResolver embeddedValueResolver;
    
    @Setter
    private ClientConfig clientConfig;
    @Setter
    private ScheduleSocketClient socketClient;
    @Setter
    private ExecutorService executorService;
    
    private ApplicationContext applicationContext;
    
    private final ScheduleClient scheduleClient = new ScheduleClient();
    
    private final Map<String, StaticTaskRegistryConfig> unregisteredStaticTasks = Collections.synchronizedMap(new LinkedHashMap<>(16));
    private final Set<Class<?>> nonAnnotatedClasses = Collections.newSetFromMap(new ConcurrentHashMap<>(64));
    
    /**
     * Setting an {@link ApplicationContext} is optional: If set, registered
     * tasks will be activated in the {@link ContextRefreshedEvent} phase;
     * if not set, it will happen at {@link #afterSingletonsInstantiated} time.
     */
    @Override
    public void setApplicationContext(@NonNull ApplicationContext applicationContext) {
        this.applicationContext = applicationContext;
    }
    
    @Override
    public void afterSingletonsInstantiated() {
        // Remove resolved singleton classes from cache
        this.nonAnnotatedClasses.clear();
        
        if (this.applicationContext == null) {
            // Not running in an ApplicationContext -> register tasks early...
            finishRegistration();
        }
    }
    
    @Override
    public void onApplicationEvent(ContextRefreshedEvent event) {
        if (event.getApplicationContext() == this.applicationContext) {
            // Running in an ApplicationContext -> register tasks this late...
            // giving other ContextRefreshedEvent listeners a chance to perform
            // their work at the same time (e.g. Spring Batch's job registration).
            finishRegistration();
        }
    }
    
    @FunctionalInterface
    private interface MethodTaskAnnotationProcessor<A extends Annotation> {
        void processAnnotationMethod(Object bean, Method method, A annotation);
    }
    
    @Override
    public Object postProcessAfterInitialization(@NonNull Object bean, @NonNull String beanName) {
        if (bean instanceof AopInfrastructureBean || bean instanceof TaskScheduler ||
                    bean instanceof ScheduledExecutorService) {
            // Ignore AOP infrastructure such as scoped proxies.
            return bean;
        }
        
        Class<?> targetClass = AopProxyUtils.ultimateTargetClass(bean);
        scanClassTargetTaskAnnotation(StaticTask.class, targetClass, bean, this::processStaticTaskAnnotationMethod);
        return bean;
    }
    
    /**
     * 为之后拓展其他任务类型抽象出方法？
     *
     * @param annotationClass 注解类型
     * @param targetClass     bean class
     * @param bean            bean Object
     * @param processor       注解处理器
     * @param <A>             注解类型
     */
    private <A extends Annotation>
    void scanClassTargetTaskAnnotation(Class<A> annotationClass, Class<?> targetClass, Object bean,
                                       MethodTaskAnnotationProcessor<A> processor) {
        if (!this.nonAnnotatedClasses.contains(targetClass) &&
                    AnnotationUtils.isCandidateClass(targetClass, Collections.singletonList(annotationClass))) {
            Map<Method, A> annotatedMethods = MethodIntrospector.selectMethods(
                    targetClass,
                    (MethodIntrospector.MetadataLookup<A>) method -> method.getAnnotation(annotationClass)
            );
            if (annotatedMethods.isEmpty()) {
                this.nonAnnotatedClasses.add(targetClass);
                if (log.isTraceEnabled()) {
                    log.trace("No @StaticTask annotations found on bean class: " + targetClass);
                }
            } else {
                // Non-empty set of methods
                annotatedMethods.forEach(
                        (method, scheduledAnnotation) ->
                                processor.processAnnotationMethod(bean, method, scheduledAnnotation));
                if (log.isTraceEnabled()) {
                    log.trace(annotatedMethods.size() + " @StaticTask methods processed on bean '" + targetClass +
                                      "': " + annotatedMethods);
                }
            }
        }
    }
    
    protected synchronized void processStaticTaskAnnotationMethod(Object bean, Method method, StaticTask annotation) {
        final ContextualMethodInjectionTaskRunnable runnable = new ContextualMethodInjectionTaskRunnable(bean, method);
        final StaticTaskRegistry registry = new StaticTaskRegistry();
        registry.setTaskDesc(embeddedValueResolver.resolveStringValue(annotation.description()));
        registry.setTaskName(embeddedValueResolver.resolveStringValue(annotation.name()));
        if (!StringUtils.hasText(registry.getTaskName())) {
            registry.setTaskName(method.toString());
        }
        registry.setTriggerExpression(embeddedValueResolver.resolveStringValue(annotation.expression()));
        final StaticTaskRegistryConfig config = new StaticTaskRegistryConfig();
        config.setRegistry(registry);
        config.setRunnable(runnable);
        final StaticTaskRegistryConfig exists = unregisteredStaticTasks.put(registry.getTaskName(), config);
        if (exists != null) {
            final Method existsMethod = ((ContextualMethodInjectionTaskRunnable) exists.getRunnable()).getMethod();
            throw new IllegalStateException("方法[" + method + "]和方法[" + existsMethod + "]任务名称[" + annotation.name() + "]重复！");
        }
    }
    
    
    private synchronized void finishRegistration() {
        scheduleClient.setClientConfig(clientConfig);
        scheduleClient.setSocketClient(socketClient);
        scheduleClient.setExecutorService(executorService);
        scheduleClient.setStaticTasks(new ArrayList<>(unregisteredStaticTasks.values()));
        unregisteredStaticTasks.clear();
        scheduleClient.start();
    }
    
    @Override
    public synchronized void destroy() {
        scheduleClient.shutdown();
        scheduleClient.getStaticTasks().clear();
    }
    
    @Override
    public void postProcessBeforeDestruction(@NonNull Object bean, @NonNull String beanName) throws BeansException {
        // 应该实现取消掉销毁掉Bean对应的任务， 这里schedule-center机制不同没有办法控制到这么细节
    }
    
    @Override
    public boolean requiresDestruction(@NonNull Object bean) {
        return false;
    }
    
    @Override
    public void postProcessMergedBeanDefinition(@NonNull RootBeanDefinition beanDefinition, @NonNull Class<?> beanType, @NonNull String beanName) {
        // do nothing
    }
    
    @RequiredArgsConstructor
    public static class NoTrigger implements Trigger {
        final String expression;
        
        @Override
        public Date nextExecutionTime(@NonNull TriggerContext triggerContext) {
            return null;
        }
        
        @Override
        public String toString() {
            return "expression [" + expression + "] can not make a trigger";
        }
    }
    
    @Override
    @NonNull
    public Set<ScheduledTask> getScheduledTasks() {
        final ScheduledTaskRegistrar scheduledTaskRegistrar = new ScheduledTaskRegistrar();
        return scheduleClient.getStaticTasks().stream().map(registry -> {
                    try {
                        final CronTrigger trigger = new CronTrigger(registry.getRegistry().getTriggerExpression());
                        return new ScheduleClientStaticTask(registry, trigger);
                    } catch (Exception e) {
                        return new ScheduleClientStaticTask(registry, new NoTrigger(registry.getRegistry().getTriggerExpression()));
                    }
                })
                       .map(scheduledTaskRegistrar::scheduleTriggerTask)
                       .collect(Collectors.toSet());
    }
}
